package com.blm.handler;

import com.blm.repository.UserRepository;
import com.blm.util.JwtUtil;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.http.server.ServerHttpRequest;
import org.springframework.http.server.ServerHttpResponse;
import org.springframework.stereotype.Component;
import org.springframework.web.socket.WebSocketHandler;
import org.springframework.web.socket.server.HandshakeInterceptor;
import org.springframework.web.util.UriComponentsBuilder;

import java.util.Map;

@Component
public class WebSocketAuthInterceptor implements HandshakeInterceptor {
    private static final Logger logger = LoggerFactory.getLogger(WebSocketAuthInterceptor.class);
    
    @Autowired
    private JwtUtil jwtUtil;
    
    @Autowired
    private UserRepository userRepository;
    
    @Override
    public boolean beforeHandshake(ServerHttpRequest request, ServerHttpResponse response, WebSocketHandler wsHandler, Map<String, Object> attributes) throws Exception {
        try {
            // Extract token from query parameters
            String token = extractTokenFromRequest(request);
            
            if (token != null && !jwtUtil.isTokenExpired(token)) {
                // Validate token and extract username
                String username = jwtUtil.extractUsername(token);
                if (username == null) {
                    logger.warn("WebSocket connection rejected: Invalid token");
                    return false;
                }
                Long userId = userRepository.getUserIdByUsername(username);
                if (userId != null) {
                    // Store user ID in WebSocket session attributes
                    attributes.put("userId", userId);
                    return true;
                }
            }
            
            logger.warn("WebSocket connection rejected: Invalid or missing token");
            return false;
        } catch (Exception e) {
            logger.error("Error during WebSocket handshake authentication", e);
            return false;
        }
    }

    @Override
    public void afterHandshake(ServerHttpRequest request, ServerHttpResponse response, WebSocketHandler wsHandler, Exception exception) {
        // No action needed after handshake
    }
    
    private String extractTokenFromRequest(ServerHttpRequest request) {
        // Parse URL and query parameters
        Map<String, String> queryParams = UriComponentsBuilder
                .fromUri(request.getURI())
                .build()
                .getQueryParams()
                .toSingleValueMap();

        return queryParams.get("token");
    }
}
